package tcp_server

import (
	"context"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

/*
TCP 服务器的基础框架

从接受到处理的完整流程：
1. tcpKeepAliveListener.Accept()
   └─ 接受新的 TCP 连接

2. TcpServer.Serve()
   └─ 创建 conn 对象
   └─ 启动 goroutine: go c.serve(ctx)

3. conn.serve()
   └─ 设置异常恢复
   └─ 记录客户端地址
   └─ 丰富上下文信息
   └─ 调用 c.server.Handler.ServeTCP(ctx, c.rwc)

4. 业务处理器 (如TCP代理中间件)
   └─ 执行中间件链
   └─ 负载均衡
   └─ 反向代理到后端

*/

var (
	ErrServerClosed = errors.New("tcp: Server closed")
	ErrAbortHandler = errors.New("tcp: abort TCPHandler")
	// 全局上下文键
	ServerContextKey    = &contextKey{"tcp-server"}
	LocalAddrContextKey = &contextKey{"local-addr"}
)

type onceCloseListener struct {
	net.Listener
	once     sync.Once
	closeErr error
}

func (oc *onceCloseListener) Close() error {
	oc.once.Do(oc.close) // 确保只执行一次关闭
	return oc.closeErr
}

func (oc *onceCloseListener) close() {
	oc.closeErr = oc.Listener.Close()
}

type TCPHandler interface {
	ServeTCP(ctx context.Context, conn net.Conn)
}

type TcpServer struct {
	Addr    string     // 监听地址
	Handler TCPHandler // 业务处理器
	err     error
	BaseCtx context.Context // 基础上下文

	// 超时控制
	WriteTimeout     time.Duration
	ReadTimeout      time.Duration
	KeepAliveTimeout time.Duration

	// 状态管理
	mu         sync.Mutex
	inShutdown int32
	doneChan   chan struct{}
	l          *onceCloseListener
}

// 原子操作状态检查
func (s *TcpServer) shuttingDown() bool {
	return atomic.LoadInt32(&s.inShutdown) != 0
}

// 启动服务器
func (srv *TcpServer) ListenAndServe() error {
	if srv.shuttingDown() {
		return ErrServerClosed
	}
	if srv.doneChan == nil {
		srv.doneChan = make(chan struct{})
	}
	addr := srv.Addr
	if addr == "" {
		return errors.New("need addr")
	}
	ln, err := net.Listen("tcp", addr) // 创建监听器
	if err != nil {
		return err
	}
	return srv.Serve(tcpKeepAliveListener{
		ln.(*net.TCPListener)})
}

// 优雅关闭
func (srv *TcpServer) Close() error {
	atomic.StoreInt32(&srv.inShutdown, 1) // 设置关闭状态
	close(srv.doneChan)                   // 关闭 channel，通知所有等待的 goroutine
	srv.l.Close()                         // 关闭监听器，停止接受新连接
	return nil
}

// 服务循环
func (srv *TcpServer) Serve(l net.Listener) error {
	srv.l = &onceCloseListener{Listener: l}
	defer srv.l.Close() //执行listener关闭
	if srv.BaseCtx == nil {
		srv.BaseCtx = context.Background()
	}
	baseCtx := srv.BaseCtx
	ctx := context.WithValue(baseCtx, ServerContextKey, srv)
	for {
		rw, e := l.Accept() // 监听客户端连接，代理服务器作为服务器接受新连接 (这里完成第一次握手：客户端 ↔ 代理服务器)
		if e != nil {
			select {
			case <-srv.getDoneChan(): // 检查是否正在关闭
				return ErrServerClosed
			default:
			}
			fmt.Printf("accept fail, err: %v\n", e)
			continue
		}
		c := srv.newConn(rw) // 创建连接对象
		go c.serve(ctx)      // 启动新的 goroutine 处理连接
	}
}

// 连接包装器
func (srv *TcpServer) newConn(rwc net.Conn) *conn {
	c := &conn{
		server: srv,
		rwc:    rwc,
	}
	// 设置超时参数
	if d := c.server.ReadTimeout; d != 0 {
		c.rwc.SetReadDeadline(time.Now().Add(d))
	}
	if d := c.server.WriteTimeout; d != 0 {
		c.rwc.SetWriteDeadline(time.Now().Add(d))
	}
	if d := c.server.KeepAliveTimeout; d != 0 {
		if tcpConn, ok := c.rwc.(*net.TCPConn); ok {
			tcpConn.SetKeepAlive(true)
			tcpConn.SetKeepAlivePeriod(d)
		}
	}
	return c
}

// doneChan 模式
func (s *TcpServer) getDoneChan() <-chan struct{} {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.doneChan == nil {
		s.doneChan = make(chan struct{})
	}
	return s.doneChan
}

func ListenAndServe(addr string, handler TCPHandler) error {
	server := &TcpServer{Addr: addr, Handler: handler, doneChan: make(chan struct{})}
	return server.ListenAndServe()
}
